"""
实验1：结构化提示词与输出
学生需要实现 classify_text 函数，使用 Pydantic 模型返回结构化的文本分类结果
"""
from typing import List
from pydantic import BaseModel, Field
import httpx
import json


class TextClassification(BaseModel):
    """
    文本分类结果的数据模型
    """
    category: str = Field(
        ...,
        description="文本分类类别，必须是以下之一：'新闻', '技术', '体育', '娱乐', '财经'"
    )
    confidence_score: float = Field(
        ...,
        ge=0.0,
        le=1.0,
        description="分类置信度，范围0.0-1.0"
    )
    keywords: List[str] = Field(
        ...,
        min_length=1,
        max_length=5,
        description="从文本中提取的1-5个关键词"
    )


def classify_text(text) -> TextClassification:
    import requests

    """
    对输入文本进行分类，返回结构化的分类结果
    
    参数:
        text: 待分类的文本内容
    
    返回:
        TextClassification 实例，包含分类类别、置信度和关键词
    
    实现要求:
        1. 使用 Ollama API 调用 qwen3:8b 模型
        2. 设计结构化 Prompt，要求模型输出 JSON 格式
        3. 解析模型输出并验证为 TextClassification 模型
        4. category 必须是预定义的5个类别之一
        5. confidence_score 必须在 0-1 范围内
        6. keywords 列表长度为 1-5
    
    提示:
        - 在 Prompt 中明确指定输出格式和有效类别
        - 可以使用 Few-Shot 示例提高输出稳定性
        - 使用 Pydantic 的自动验证确保数据有效性
    """
    # 构建结构化Prompt，要求模型输出JSON格式
    prompt = f"""请对以下文本进行分类，输出格式为JSON，必须包含以下字段：
                1. category: 文本分类类别，必须是以下之一：'新闻', '技术', '体育', '娱乐', '财经'
                2. confidence_score: 分类置信度，范围0.0-1.0的浮点数
                3. keywords: 从文本中提取的1-5个关键词，以数组形式返回

                文本内容：{text}

                请严格按照以下格式输出JSON：
                {{
                "category": "[类别]",
                "confidence_score": [置信度],
                "keywords": ["关键词1", "关键词2", ...]
                }}

                示例输出：
                {{
                "category": "技术",
                "confidence_score": 0.95,
                "keywords": ["人工智能", "医疗诊断", "研究"]
                }}"""

    # 调用Ollama API
    url = "http://localhost:11434/api/generate"
    payload = {
        "model": "qwen3:8b",  # 使用qwen3:8b模型
        "prompt": prompt,
        "format": "json",
        "stream": False
    }

    try:
        # 尝试发送请求到Ollama API
        with httpx.Client() as client:
            response = client.post(url, json=payload)
            response.raise_for_status()  # 检查请求是否成功
        # response = requests.post(url, json=payload)
        # response.raise_for_status()  # 检查请求是否成功

            # 解析响应中的JSON字符串
            result = response.json()
            response_text = result.get("response", "")

            # 尝试解析响应文本中的JSON
            try:
                classification_data = json.loads(response_text)
            except json.JSONDecodeError:
                # 如果响应不是有效的JSON，尝试提取JSON部分
                start = response_text.find("{")
                end = response_text.rfind("}") + 1
                if start != -1 and end != 0:
                    json_str = response_text[start:end]
                    classification_data = json.loads(json_str)
                else:
                    raise ValueError("无法从响应中提取有效的JSON格式")

            # 使用TextClassification.model_validate()创建实例
            classification = TextClassification.model_validate(classification_data)
            print(f"使用Ollama API分类结果: {classification}")

            # 返回验证后的Pydantic模型实例
            return classification



    except Exception as e:
        # 其他异常情况下的处理
        print(f"文本分类过程中发生错误: {str(e)}，使用默认分类")


    raise NotImplementedError("请实现 classify_text 函数")


# 测试代码（可选，用于学生本地调试）
if __name__ == "__main__":
    # 测试示例
    test_texts = [
        "OpenAI发布GPT-5，性能提升10倍",
        "中国队在巴黎奥运会夺得金牌",
        "A股市场今日大涨，沪指突破3000点"
    ]

    for text in test_texts:
        try:
            result = classify_text(text)
            print(f"\n文本: {text}")
            print(f"分类: {result.category}")
            print(f"置信度: {result.confidence_score}")
            print(f"关键词: {result.keywords}")
        except Exception as e:
            print(f"错误: {e}")
